package cn.trigram.structure;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;
import java.util.ListIterator;
import java.util.Objects;
import java.util.function.Consumer;
import org.apache.commons.collections4.CollectionUtils;

/**
 * 多叉树的工具类。以下例的树为例子：
 * <pre>
 *       A
 *   ↙   ↓   ↘
 *  B    C    D
 *     ↙   ↘  ↓
 *    E     F G
 * </pre>
 * 多叉树的中序遍历基本没有意义，所以没有中序遍历
 *
 * @author 一日看尽长安花
 * @date 2022/01/12 22:03:07
 */
public class MultiWayTreeUtil {

  /**
   * 数组转树
   *
   * @param root
   *     根节点
   * @param nodes
   *     节点
   *
   * @return {@link TreeNode}
   */
  public static TreeNode toTree(TreeNode root, List<? extends TreeNode> nodes) {

    if (CollectionUtils.isEmpty(nodes)) {
      return root;
    }
    /*采用前序遍历的方式迭代填充子节点*/
    Deque<TreeNode> stack = new ArrayDeque<>();
    stack.addFirst(root);
    while (!stack.isEmpty()) {
      TreeNode midNode = stack.removeFirst();
      if (CollectionUtils.isEmpty(nodes)) {
        break;
      }
      List<TreeNode> children = new ArrayList<>();
      for (ListIterator<? extends TreeNode> iterator = nodes.listIterator(); iterator.hasNext(); ) {
        TreeNode node = iterator.next();
        if (midNode.isParent(node)) {
          children.add(node);
          stack.addFirst(node);
          iterator.remove();
        }
      }
      if (CollectionUtils.isNotEmpty(children)) {
        midNode.setChildren(children);
      }
    }
    return root;
  }

  /**
   * 默认的根节点是null的情况
   *
   * @param nodes
   *     节点
   *
   * @return {@link TreeNode}
   */
  public static TreeNode toTree(List<? extends TreeNode> nodes) {

    return toTree(new TreeNode() {

      @Override
      public boolean hasParent() {

        return false;
      }

      @Override
      public boolean equalsTo(TreeNode node) {

        return Objects.isNull(node) || this.equals(node);
      }

      @Override
      public boolean isParent(TreeNode node) {

        return !node.hasParent();
      }

      @Override
      public String toString() {

        return "";
      }
    }, nodes);
  }


  /**
   * 多叉树的前序非递归遍历
   * <p>
   * 遍历结果：A->B->C->E->F->D->G.
   */
  public static void preOrderNonRecursive(
      TreeNode root,
      Consumer<TreeNode> consumer
  ) {

    if (Objects.isNull(root)) {
      return;
    }
    Deque<TreeNode> stack = new ArrayDeque<>();
    stack.addFirst(root);
    while (!stack.isEmpty()) {
      TreeNode node = stack.removeFirst();
      consumer.accept(node);
      List<TreeNode> children = node.getChildren();
      if (CollectionUtils.isEmpty(children)) {
        continue;
      }
      /*遍历时先左后右，入栈就要先右后左，所以这里逆序遍历*/
      for (
          ListIterator<TreeNode> iterator = children.listIterator(children.size());
          iterator.hasPrevious(); ) {
        TreeNode previous = iterator.previous();
        stack.addFirst(previous);
      }
    }
  }

  /**
   * 多叉树的后序非递归遍历。左->右->中
   * <p>
   * 遍历结果：B->E->F->C->G->D->A.
   *
   * @param root
   *     根
   */
  public static void postOrderNonRecursive(TreeNode root, Consumer<TreeNode> consumer) {

    if (Objects.isNull(root)) {
      return;
    }

    Deque<TreeNode> stack = new ArrayDeque<>();
    stack.addFirst(root);
    /*指向上一次出栈的节点*/
    TreeNode pointer = null;

    while (!stack.isEmpty()) {
      TreeNode       node         = stack.peekFirst();
      List<TreeNode> nodeChildren = node.getChildren();
      if (CollectionUtils.isEmpty(nodeChildren) || CollectionUtils.containsAny(
          nodeChildren, pointer)) {
        // 1、如果是叶子节点，则表示已经可以出栈
        // 2、上一次访问的节点是此节点的子节点时，表示当前节点的子节点已经全部出栈了，此时当前节点为父节点应该出栈访问
        consumer.accept(node);
        pointer = node;
        stack.removeFirst();
      } else {
        if (CollectionUtils.isNotEmpty(nodeChildren)) {

          for (
              ListIterator<TreeNode> listIterator = nodeChildren.listIterator(
                  nodeChildren.size()); listIterator.hasPrevious(); ) {
            stack.addFirst(listIterator.previous());
          }
        }
      }
    }
  }

  /**
   * 多叉树的后序非递归遍历的更简单版本。左->右->中
   * <p>
   * 遍历结果：B->E->F->C->G->D->A.
   *
   * @param root
   *     根
   */
  public static void postOrderSimpleNonRecursive(TreeNode root, Consumer<TreeNode> consumer) {

    if (Objects.isNull(root)) {
      return;
    }
    Deque<TreeNode> resultDeque = new ArrayDeque<>();
    Deque<TreeNode> stack       = new ArrayDeque<>();
    stack.push(root);

    while (!stack.isEmpty()) {
      TreeNode node = stack.pop();
      if (CollectionUtils.isNotEmpty(node.getChildren())) {
        for (TreeNode tmpNode : node.getChildren()) {
          stack.push(tmpNode);
        }
      }
      resultDeque.addLast(node);
    }
    while (!resultDeque.isEmpty()) {
      consumer.accept(resultDeque.removeLast());
    }
  }


  /**
   * 层序遍历非递归
   * <p>
   * 遍历结果: A->B->C->D->E->F->G.
   */
  public static void sequenceNonRecursive(TreeNode root, Consumer<TreeNode> consumer) {

    if (Objects.isNull(root)) {
      return;
    }
    /*迭代队列*/
    Deque<TreeNode> iteratorDeque = new ArrayDeque<>();
    /*消费队列*/
    Deque<TreeNode> consumerDeque = new ArrayDeque<>();
    consumerDeque.addFirst(root);
    while (!consumerDeque.isEmpty()) {
      while (!consumerDeque.isEmpty()) {
        TreeNode node = consumerDeque.removeFirst();
        consumer.accept(node);
        iteratorDeque.addLast(node);
      }
      while (!iteratorDeque.isEmpty()) {
        TreeNode       node     = iteratorDeque.removeFirst();
        List<TreeNode> children = node.getChildren();
        if (CollectionUtils.isEmpty(children)) {
          continue;
        }
        for (TreeNode child : children) {
          consumerDeque.addLast(child);
        }
      }
    }
  }

  /**
   * 从一颗树上遍历所有分支，修剪掉给定的端节点的子节点。
   * <p>
   * 很像给树修剪树枝。
   *
   * @param root
   *     根
   */
  public static TreeNode trimTreeBranch(TreeNode root, List<? extends TreeNode> endNodes) {

    if (Objects.isNull(root)) {
      return null;
    }
    /**/
    Deque<TreeNode> retainDeque = new ArrayDeque<>();
    Deque<TreeNode> stack       = new ArrayDeque<>();
    stack.addFirst(root);
    /*指向上一次出栈的节点*/
    TreeNode pointer = null;
    /*上一个保留节点的指针*/
    TreeNode retainPointer = null;

    while (!stack.isEmpty()) {
      TreeNode       node         = stack.peekFirst();
      List<TreeNode> nodeChildren = node.getChildren();
      if (CollectionUtils.isEmpty(nodeChildren) || CollectionUtils.containsAny(
          nodeChildren, pointer)) {
        // 条件：A || B
        // 1、如果是叶子节点，A则成立，子节点出栈
        // 2、如果此时当前节点为父节点，A不成立，但B由于上一次访问的节点是当前节点的子节点所以成立，
        //    此时当前节点为父节点应该出栈访问
        if (endNodes.contains(node) && !retainDeque.contains(node)) {
          retainPointer = node;
          retainDeque.addLast(node.clone());
        }
        if (
            CollectionUtils.isNotEmpty(nodeChildren) &&
            Objects.equals(pointer, nodeChildren.get(nodeChildren.size() - 1))
        ) {
          /*当前条件表示当前栈顶出栈的是父节点，需要判断子节点中是否有在保留队列中*/
          if (Objects.nonNull(retainPointer) && !retainDeque.contains(node)) {
            /*判断当前节点是否在保留节点队列中用于确保父子节点同时在保留节点的边界情况下不重复*/
            retainPointer = node;
            retainDeque.addLast(node.clone());
          }
        }
        pointer = node;
        stack.removeFirst();
      } else {
        if (CollectionUtils.isNotEmpty(nodeChildren)) {

          for (
              ListIterator<TreeNode> listIterator = nodeChildren.listIterator(
                  nodeChildren.size()); listIterator.hasPrevious(); ) {
            stack.addFirst(listIterator.previous());
          }
        }
      }
    }

    TreeNode copyRoot = retainDeque.removeLast();
    return toTree(copyRoot, new ArrayList<>(retainDeque));
  }

}
